import json
import logging
import time
from datetime import datetime
from pathlib import Path
from zhipuai import ZhipuAI
import re
import traceback
import os
import fitz
import pytesseract
from PIL import Image
import io

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('contract_analysis.log', encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)

# 设置Tesseract OCR路径
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

class ContractAnalyzer:
    def __init__(self, api_key: str = "1a35c255bdd44acb8a1686ff7c3948d2.SdLssvRgDvCZJCog"):
        """
        初始化智谱清言客户端
        
        Args:
            api_key: 智谱清言API密钥 
        """
        self.api_key = api_key
        self.client = ZhipuAI(api_key=api_key)
        self.logger = logging.getLogger(__name__)
        self.logger.info("智谱清言客户端初始化成功")
        
        # 定义合同分析模板
        self.template = {
            "合同编号": None,
            "合同名称": None,
            "合同类型": None,
            "申请日期": None,
            "附件上传": None,
            "签订日期": None,
            "甲方信息": {
                "名称": None,
                "地址": None,
                "法定代表人": None,
                "联系方式": None
            },
            "乙方信息": {
                "名称": None,
                "地址": None,
                "法定代表人": None,
                "联系方式": None
            },
            "财务信息": {
                "收款金额（小写）": None,
                "收款金额（大写）": None,
                "税率": None,
                "乙方收款账号": None
            }
        }
        
    def _extract_text_from_image(self, image_data: bytes) -> str:
        """
        从图像中提取文本
        Args:
            image_data: 图像数据
        Returns:
            str: 提取的文本
        """
        try:
            # 将图像数据转换为PIL Image
            image = Image.open(io.BytesIO(image_data))
            
            # 转换为灰度图像
            if image.mode != 'L':
                image = image.convert('L')
            
            # 使用tesseract进行OCR（使用中文和英文）
            text = pytesseract.image_to_string(image, lang='chi_sim+eng')
            
            return text.strip()
            
        except Exception as e:
            self.logger.error(f"OCR处理失败: {str(e)}")
            return ""
        
    def _extract_pdf_text(self, file_path: str) -> str:
        """
        从PDF文件中提取文本
        Args:
            file_path: PDF文件路径
        Returns:
            str: 提取的文本内容
        """
        try:
            self.logger.info(f"开始读取PDF文件: {file_path}")
            start_time = time.time()
            
            # 打开PDF文件
            doc = fitz.open(file_path)
            self.logger.info(f"PDF文件页数: {len(doc)}")
            
            text = ""
            # 遍历每一页并提取文本
            for page_num, page in enumerate(doc, 1):
                self.logger.info(f"正在处理第 {page_num} 页...")
                
                # 获取页面图像
                pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))  # 2x放大以提高OCR质量
                img_data = pix.tobytes()
                
                # 使用OCR提取文本
                page_text = self._extract_text_from_image(img_data)
                
                if page_text:
                    self.logger.info(f"第 {page_num} 页文本长度: {len(page_text)}")
                    # 打印前100个字符用于调试
                    self.logger.info(f"第 {page_num} 页文本示例: {page_text[:100]}")
                    text += page_text + "\n"  # 添加换行符分隔页面
                else:
                    self.logger.warning(f"第 {page_num} 页未提取到有效文本")
            
            # 关闭PDF文件
            doc.close()
            
            # 清理最终文本
            text = text.strip()
            
            if not text:
                raise ValueError("未能从PDF中提取到文本")
            
            # 打印提取到的文本示例
            self.logger.info(f"提取的文本示例:\n{text[:500]}")
            self.logger.info(f"PDF文本提取成功，总长度: {len(text)} 字符")
            self.logger.info(f"文本提取耗时: {time.time() - start_time:.2f}秒")
            
            return text
            
        except Exception as e:
            self.logger.error(f"提取PDF文本时发生错误: {str(e)}")
            raise ValueError(f"PDF文本提取失败: {str(e)}")
            
    def process_file(self, file_path: str = "F:/下载/合同2.pdf") -> dict:
        """
        处理合同文件并分析内容
        Args:
            file_path: 合同文件路径
        Returns:
            dict: 分析结果
        """
        try:
            start_time = time.time()
            
            # 检查文件路径
            self.logger.info(f"文件路径: {file_path}")
            self.logger.info(f"文件是否存在: {os.path.exists(file_path)}")
            
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"文件不存在: {file_path}")
            
            # 使用PyMuPDF提取文本
            self.logger.info("开始提取PDF文本...")
            text_content = self._extract_pdf_text(file_path)
            
            if not text_content:
                raise ValueError("文件内容为空")
            
            # 分析合同内容
            self.logger.info("开始分析合同内容...")
            result = self._analyze_contract(text_content)
            
            # 计算执行时间
            end_time = time.time()
            total_time = end_time - start_time
            
            self.logger.info(f"处理完成，总耗时: {total_time:.2f}秒")
            
            return {
                "status": "success",
                "result": result,
                "execution_time": {
                    "total": total_time
                }
            }
            
        except FileNotFoundError as e:
            error_info = self._get_error_info()
            self.logger.error(f"文件不存在: {file_path}\n{error_info}")
            return {
                "status": "error",
                "error_type": "FileNotFoundError",
                "message": f"文件不存在: {file_path}",
                "traceback": error_info
            }
        except ValueError as e:
            error_info = self._get_error_info()
            self.logger.error(f"值错误: {str(e)}\n{error_info}")
            return {
                "status": "error",
                "error_type": "ValueError",
                "message": str(e),
                "traceback": error_info
            }
        except Exception as e:
            error_info = self._get_error_info()
            self.logger.error(f"处理文件时发生错误: {str(e)}\n{error_info}")
            return {
                "status": "error",
                "error_type": type(e).__name__,
                "message": str(e),
                "traceback": error_info
            }
            
    def _analyze_contract(self, content: str) -> dict:
        """
        分析合同内容
        Args:
            content: 合同文本内容
        Returns:
            dict: 分析结果
        """
        try:
            self.logger.info("开始分析合同内容...")
            start_time = time.time()
            
            if not content or len(content.strip()) == 0:
                raise ValueError("合同内容为空")
            
            self.logger.info(f"合同内容长度: {len(content)} 字符")
            
            # 构建分析消息
            messages = [
                {
                    "role": "system",
                    "content": f"""你是一个专业的合同分析助手。请仔细分析合同内容，提取关键信息并按照以下JSON模板格式返回：

{json.dumps(self.template, ensure_ascii=False, indent=4)}

注意：
1. 只返回JSON数据，不要添加任何其他内容
2. 如果某个字段在合同中未找到，请保持为null
3. 金额格式：小写金额需要带¥符号和小数点，大写金额需要带"人民币"和"整"字
4. 所有日期格式统一为"YYYY年MM月DD日"
5. 严格按照模板格式返回，不要改变字段名称和层级结构"""
                },
                {
                    "role": "user",
                    "content": f"请分析以下合同内容，并提取关键信息：\n\n{content}"
                }
            ]
            
            # 调用chat API进行分析
            response = self.client.chat.completions.create(
                model="glm-4-long",  # 使用支持长文本的模型
                messages=messages,
                temperature=0.1,
                top_p=0.1,
                max_tokens=4000
            )
            
            if not response or not response.choices:
                raise Exception("未收到有效的响应")
                
            result = response.choices[0].message.content.strip()
            self.logger.info(f"API返回结果: {result[:200]}")
            
            try:
                # 尝试解析JSON
                # 首先尝试直接解析
                try:
                    parsed_result = json.loads(result)
                except json.JSONDecodeError:
                    # 如果直接解析失败，尝试提取JSON部分
                    json_pattern = r'\{[\s\S]*\}'
                    json_match = re.search(json_pattern, result)
                    if json_match:
                        json_str = json_match.group()
                        parsed_result = json.loads(json_str)
                    else:
                        raise Exception("无法从响应中提取有效的JSON")
                
                self.logger.info(f"合同分析成功，用时: {time.time() - start_time:.2f} 秒")
                return parsed_result
                
            except Exception as e:
                self.logger.error(f"JSON解析失败: {str(e)}")
                self.logger.error(f"原始响应: {result}")
                raise Exception(f"JSON解析失败: {str(e)}")
            
        except Exception as e:
            self.logger.error(f"分析合同内容时发生错误: {str(e)}")
            raise Exception(f"分析合同内容失败: {str(e)}")
            
    def _extract_json_from_markdown(self, text: str) -> str:
        """从markdown文本中提取JSON"""
        # 尝试匹配代码块中的JSON
        code_block_pattern = r"```(?:json)?\s*([\s\S]*?)\s*```"
        matches = re.findall(code_block_pattern, text)
        if matches:
            return matches[0].strip()
            
        # 尝试匹配内联代码中的JSON
        inline_pattern = r"`([\s\S]*?)`"
        matches = re.findall(inline_pattern, text)
        if matches:
            return matches[0].strip()
            
        # 如果没有markdown格式，直接返回原文本
        return text.strip()
            
    def _format_error_response(self, error_msg: str, start_time: float) -> str:
        """格式化错误响应"""
        return json.dumps({
            "status": "error",
            "message": error_msg,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "execution_time": time.time() - start_time
        }, ensure_ascii=False, indent=2)

    def _get_error_info(self) -> str:
        """获取详细的错误信息，包括堆栈跟踪"""
        return traceback.format_exc()

def test_analysis():
    """测试合同分析功能"""
    try:
        # 初始化分析器
        analyzer = ContractAnalyzer(api_key="1a35c255bdd44acb8a1686ff7c3948d2.SdLssvRgDvCZJCog")
        
        # 处理文件
        result = analyzer.process_file()
        
        # 打印结果
        print("\n执行结果:")
        if isinstance(result, dict):
            if result["status"] == "error":
                print(f"状态: 失败\n")
                print(f"错误信息: {result['message']}\n")
                if "execution_time" in result:
                    print_execution_time(result["execution_time"])
            else:
                print(f"状态: 成功\n")
                print("分析结果:")
                print(json.dumps(result["result"], ensure_ascii=False, indent=2))
                print()
                if "execution_time" in result:
                    print_execution_time(result["execution_time"])
        else:
            print(f"未知的返回结果格式: {result}")
            
    except Exception as e:
        print(f"\n执行过程中发生错误: {str(e)}")

def print_execution_time(exec_time):
    """打印执行时间统计"""
    if isinstance(exec_time, dict):
        print(f"\n执行时间统计:")
        for key, value in exec_time.items():
            print(f"{key}: {value:.2f}秒")
    else:
        print(f"\n总耗时: {exec_time:.2f}秒")

if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )
    
    # 测试分析功能
    test_analysis() 